import os
import re
from datetime import datetime

import yaml
from langchain_milvus import Milvus
from langchain_ollama import OllamaEmbeddings

# 初始化嵌入模型
embedding_model = OllamaEmbeddings(model="nomic-embed-text:latest", base_url="http://192.168.7.3:11434")

# 连接 Milvus
vectorstore = Milvus(
    embedding_function=embedding_model,
    collection_name="QA",
    connection_args={"uri": "http://192.168.6.20:19530/ipp_air_general"},
    auto_id=True
)


def normalize_metadata(metadata):
    """确保所有元数据字段都转换为支持的类型"""
    for key, value in metadata.items():
        if isinstance(value, str):
            continue  # 字符串类型无需转换
        if isinstance(value, datetime):
            metadata[key] = str(value)  # 将日期转换为字符串
        elif isinstance(value, (int, float)):
            continue  # 数字类型无需转换
        else:
            metadata[key] = str(value)  # 将其他类型转换为字符串
    return metadata


def parse_markdown(md_file):
    """解析 Markdown 文件，提取元数据和 Q&A 块"""
    with open(md_file, "r", encoding="utf-8") as f:
        content = f.read()

    # 提取元数据
    metadata_match = re.search(r"---\n(.*?)\n---", content, re.DOTALL)
    if metadata_match:
        metadata_str = metadata_match.group(1)
        try:
            metadata = yaml.safe_load(metadata_str)  # 尝试解析 YAML
            if isinstance(metadata, dict):  # 确保解析结果是字典
                for metadatum in metadata:
                    print(f"{metadatum}: {metadata[metadatum]}")
        except yaml.YAMLError as e:
            print(f"YAML 解析错误: {e}")
            metadata = {}
    else:
        metadata = {}  # 如果没有找到 YAML front matter，设置为空字典

    # 提取标题
    title_match = re.search(r"^# (.+)", content, re.MULTILINE)
    if title_match:
        metadata["title"] = title_match.group(1)

    # 提取所有 Q&A
    qa_blocks = re.findall(r"## (.*?)\n(.*?)\n(?=## |$)", content, re.DOTALL)
    # 格式化metadata
    metadata = normalize_metadata(metadata)
    return metadata, qa_blocks


def store_in_milvus(md_file):
    """解析 Markdown 并存入 Milvus"""
    metadata, qa_blocks = parse_markdown(md_file)

    docs = []
    for question, answer in qa_blocks:
        text = f"Q: {question}\nA: {answer.strip()}"
        doc_metadata = metadata.copy()
        doc_metadata["title"] = question
        docs.append((text, doc_metadata))

    # 插入 Milvus
    vectorstore.add_texts(texts=[d[0] for d in docs], metadatas=[d[1] for d in docs])
    print(f"成功存入 {len(docs)} 条 Q&A")


if __name__ == "__main__":
    source_dir = r"C:\Desktop\各类知识\问答"  # 设置源目录路径
    # 遍历输入目录下的所有文件
    for filename in os.listdir(source_dir):
        if filename.lower().endswith(".md"):  # 只处理 .md 文件
            input_path = os.path.join(source_dir, filename)
            print(f"正在处理文件: {input_path}")
            store_in_milvus(input_path)  # 处理并存储到 Milvus
